from flask import Flask, request, jsonify
from langchain import PromptTemplate, LLMChain
from langchain.llms import OpenAI
import json
from flask_cors import CORS  # 导入 CORS

app = Flask(__name__)
CORS(app)  # 启用 CORS，允许所有来源访问

# 初始化 LangChain
# 设置 OpenAI API 密钥
import os
os.environ["OPENAI_API_KEY"] = "sk-fhfopxxokpblqjbivoyhwhftlbjwrgeddccbytodqipkaqum"
base_url="https://api.siliconflow.cn/v1"
# 配置 SiliconFlow 的 OpenAI 兼容 API
llm = OpenAI(
    openai_api_base=base_url,
    api_key=os.environ.get("OPENAI_API_KEY"),
    # 不是推理模型
    model_name="THUDM/glm-4-9b-chat",
    #model_name="THUDM/chatglm3-6b",
    temperature=0.7,
    top_p=0.7,
    max_tokens=1024
)
prompt = PromptTemplate(
    input_variables=["elements"],
    template="""
    通过对话的方式玩一个类似「涂鸦上帝」的元素合成游戏。
    用户输入：{elements}
    规则：
    1. 游戏开始时，用户输入{elements}
    2. 游戏会根据用户输入的{elements}随机生成新的物质
    3. 你只需要返回新的物质的名称，以及新物质的emoji
    """
)


llm_chain = LLMChain(llm=llm, prompt=prompt)

@app.route('/synthesize_elements', methods=['POST'])
def synthesize_elements():
    data = request.get_json()
    elements_str = data.get('elements')
    print(elements_str)
    if not elements_str:
        return jsonify({"error": "请提供元素数组"}), 400
    
    elements = elements_str
    if len(elements) != 2:
        return jsonify({"error": "请提供两个元素进行合成"}), 400
    
    elements = elements_str[0] +"+"+elements_str[1]
    print("elements",elements)
    # 修改此处，使用 llm_chain 来生成结果
    result = llm_chain.run(elements=elements)
    print("组装后的 prompt:", llm_chain.prompt.format(elements=elements))

    prompt2 = PromptTemplate(
        input_variables=["param_content"],
        template="""问题： {elements} 可以合成什么？
        根据以下内容选出最合适的一个回答：
        内容：{param_content}
        字数不超过10个字，不要输出其他文字、解释或多余内容。
        """
    )
    llm_chain2 = LLMChain(llm=llm, prompt=prompt2)
    print("组装后的 prompt2:", llm_chain.prompt.format(elements=result))

    result = llm_chain2.run(param_content=result)
    print("--------------------")
    print("result ",result)
    # 只取第一个合成结果
    first_result = result
    if "不知道" in first_result:
        new_element = None
    else:
        new_element = first_result
    print(new_element)
    return jsonify({"data": new_element})

if __name__ == '__main__':
    app.run(debug=True, port=5001)